import os
import torch

class Config:
    # 模型配置
    SAM_MODEL_PATH = os.environ.get('SAM_MODEL_PATH', './models/sam_vit_h_4b8939.pth')
    # 支持的模型类型: 'vit_h', 'vit_l', 'vit_b', 'mobile_sam'
    MODEL_TYPE = os.environ.get('SAM_MODEL_TYPE', 'vit_h')
    
    # 设备配置
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    # 如果有多个GPU，可指定GPU索引，如 'cuda:0', 'cuda:1' 等
    GPU_INDEX = os.environ.get('GPU_INDEX', '0')
    
    # 服务配置
    HOST = os.environ.get('HOST', '0.0.0.0')
    PORT = int(os.environ.get('PORT', 8077))
    
    # Ray Serve配置
    SERVE_NUM_REPLICAS = int(os.environ.get('SERVE_NUM_REPLICAS', '1'))
    # 每个副本使用的GPU数量
    SERVE_GPUS_PER_REPLICA = float(os.environ.get('SERVE_GPUS_PER_REPLICA', '1.0'))
    
    # 缓存配置
    ENABLE_CACHE = os.environ.get('ENABLE_CACHE', 'True').lower() == 'true'
    CACHE_SIZE = int(os.environ.get('CACHE_SIZE', '100'))
    
    # 预处理配置
    IMAGE_RESIZE = int(os.environ.get('IMAGE_RESIZE', '1024'))
    
    # 日志配置
    LOG_LEVEL = os.environ.get('LOG_LEVEL', 'INFO')

# 创建配置实例
def get_config():
    config = Config()
    # 如果有GPU可用，设置GPU索引
    if config.DEVICE == 'cuda' and torch.cuda.is_available():
        device_ids = list(range(torch.cuda.device_count()))
        if config.GPU_INDEX.isdigit() and int(config.GPU_INDEX) in device_ids:
            config.DEVICE = f'cuda:{config.GPU_INDEX}'
    return config